
#include "lin_space_d_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;                                                       

static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    uint64_t ubSize, tilingKey;
    uint32_t ubDataNumber = 20;
    uint32_t dataTypeSize = 4;
    LinSpaceDCustomTilingData tiling;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    auto coreNum = ascendcPlatform.GetCoreNum();

    auto tensorN = context->GetInputTensor(2);
    auto dataPtr = tensorN->GetData<int64_t>();
    const int64_t gettingTotalLength = dataPtr[0];
    uint32_t totalLength = static_cast<uint32_t>(gettingTotalLength);

    int start_type = context->GetInputDesc(0)->GetDataType(); 
    int end_type = context->GetInputDesc(1)->GetDataType(); 

    //预留double接口
    if (start_type == ge::DT_DOUBLE || end_type == ge::DT_DOUBLE) {
        tilingKey = 101;
        dataTypeSize = 8;
    }
    else if (
        (start_type == ge::DT_UINT8 || start_type == ge::DT_BF16 ||
         start_type == ge::DT_FLOAT16 || start_type == ge::DT_FLOAT || 
         start_type == ge::DT_INT8 || start_type == ge::DT_INT16 || start_type == ge::DT_INT32) &&
        (end_type == ge::DT_UINT8  || end_type == ge::DT_BF16 ||
         end_type == ge::DT_FLOAT16 || end_type == ge::DT_FLOAT || 
         end_type == ge::DT_INT8 || end_type == ge::DT_INT16 || end_type == ge::DT_INT32)
    ) {
        tilingKey = 100;
    }
    else {
        return ge::GRAPH_FAILED;
    }
    auto ret = context->SetTilingKey(tilingKey);

    uint32_t alignNum = BLOCK_SIZE / dataTypeSize;
    uint32_t ubBlockNum = (ubSize / BLOCK_SIZE) / ubDataNumber;
    uint32_t totalLengthAligned = (totalLength % alignNum == 0U)
        ? totalLength
        : ((totalLength + alignNum - 1) / alignNum) * alignNum;
  
    uint32_t totalBlockCount = totalLengthAligned / alignNum;  
    uint32_t formerNum = totalBlockCount % coreNum;           
    uint32_t tailNum = coreNum - formerNum;               
    uint32_t formerLength, tailLength, fTileNum, fLastTileLen, tTileNum, tLastTileLen;
    uint32_t maxPerCoreElem = alignNum * ubBlockNum;
    if(formerNum == 0) {
        formerLength = 0;
        fTileNum =0;
        fLastTileLen = 0;
    }
    else {
        formerLength = (totalBlockCount + coreNum - 1) / coreNum * alignNum;
        fTileNum = formerLength / maxPerCoreElem; 
            if ((formerLength % maxPerCoreElem == 0U) || (fTileNum == 0U)) {
            if (fTileNum == 0U) {
                fTileNum = (formerLength + alignNum - 1) / alignNum;
            }                                                                      
            if (formerLength < maxPerCoreElem) {
                uint32_t blockCount = (formerLength + alignNum - 1) / alignNum;  
                blockCount = (blockCount + BUFFER_NUM - 1) / BUFFER_NUM * BUFFER_NUM;                             
                fLastTileLen = (fTileNum == 1) ? formerLength : (formerLength % alignNum == 0 ? alignNum : formerLength % alignNum);
            } 
            else {
                fLastTileLen = formerLength - (fTileNum - 1) * maxPerCoreElem;
            }
        } 
        else {
            fTileNum++;
            fLastTileLen = formerLength - (fTileNum - 1) * maxPerCoreElem;
        }
    }

    tailLength = totalBlockCount / coreNum * alignNum;
    tTileNum = tailLength / maxPerCoreElem;  
    if ((tailLength % maxPerCoreElem == 0U) || (tTileNum == 0U)) {
        if (tTileNum == 0U) {
            tTileNum = (tailLength + alignNum - 1) / alignNum;
        }                                                                      
        if (tailLength < maxPerCoreElem) {
            uint32_t blockCount = (tailLength + alignNum - 1) / alignNum;  
            blockCount = (blockCount + BUFFER_NUM - 1) / BUFFER_NUM * BUFFER_NUM;                             
            tLastTileLen = (tTileNum == 1) ? tailLength : (tailLength % alignNum == 0 ? alignNum : tailLength % alignNum);
        } 
        else {
            tLastTileLen = tailLength - (tTileNum - 1) * maxPerCoreElem;
        }
    } 
    else {
        tTileNum++;
        tLastTileLen = tailLength - (tTileNum - 1) * maxPerCoreElem;
    }

    tiling.set_formerNum(formerNum);
    tiling.set_formerLength(formerLength);
    tiling.set_formerTileNum(fTileNum);
    tiling.set_formerLastTileLength(fLastTileLen);
    tiling.set_tailLength(tailLength);
    tiling.set_tailTileNum(tTileNum);
    tiling.set_tailLastTileLength(tLastTileLen);  
    tiling.set_totalLength(totalLength);
         
    context->SetBlockDim(coreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
} // namespace optiling


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    auto tensorData = context->GetInputTensor(2);
    if (tensorData == nullptr) return ge::GRAPH_FAILED;
    auto dataPtr = tensorData->GetData<int64_t>();
    const uint32_t num_value = static_cast<uint32_t>(dataPtr[0]);
    
    gert::Shape* y_shape = context->GetOutputShape(0); 
    if (y_shape == nullptr) return ge::GRAPH_FAILED;
    
    y_shape->SetDimNum(1);         
    y_shape->SetDim(0, num_value); 
    return ge::GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
    context->SetOutputDataType(0, ge::DT_FLOAT);
    return ge::GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class LinSpaceDCustom : public OpDef {
public:
    explicit LinSpaceDCustom(const char* name) : OpDef(name)
    {
        this->Input("start")
            .ParamType(REQUIRED)
            .DataTypeList({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_UINT8, ge::DT_BF16})
            .FormatList({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .AutoContiguous();
        this->Input("end")
            .ParamType(REQUIRED)
            .DataTypeList({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_UINT8, ge::DT_BF16})
            .FormatList({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .AutoContiguous();
        this->Input("num")
            .ParamType(REQUIRED)
            .ValueDepend(REQUIRED)
            .DataTypeList({ge::DT_INT64})
            .FormatList({ge::FORMAT_ND})
            .AutoContiguous();
        this->Output("z")
            .ParamType(REQUIRED)
            .DataTypeList({ge::DT_FLOAT})
            .FormatList({ge::FORMAT_ND});
        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore()
            .AddConfig("ascend910")	
            .AddConfig("ascend310p")	
            .AddConfig("ascend310b")	
            .AddConfig("ascend910b");
    }
};
OP_ADD(LinSpaceDCustom);
} // namespace ops